from django.db import models
from base_configs.models import EmbeddingType
import time
import uuid
from history_qa.milvus_manage import MilvusDB
from text_embedding import embedding_name2dim_dict, choose_model
from datetime import datetime


class ProjectInfo(models.Model):
    name = models.CharField(max_length=128, unique=True, verbose_name='项目名称')
    real_name = models.CharField(max_length=128, unique=True, verbose_name='实际项目名称')
    dims = models.CharField(max_length=256, blank=True, null=True, verbose_name='向量维度')
    description = models.CharField(max_length=256, verbose_name='项目描述', blank=True, null=True)
    status = models.BooleanField(choices=((1, '有效'), (0, '无效')), verbose_name='状态', default=1)
    comment = models.CharField(max_length=256, verbose_name='备注', blank=True, null=True)
    create_time = models.DateTimeField(auto_now_add=True, null=False, blank=False, verbose_name='创建时间')
    update_time = models.DateTimeField(auto_now=True, verbose_name='更新时间')

    def __str__(self):
        return self.name

    class Meta:
        verbose_name = "项目管理"
        verbose_name_plural = verbose_name


class PartitionInfo(models.Model):
    name = models.CharField(max_length=32, verbose_name='数据库名称')
    description = models.CharField(max_length=256, verbose_name='数据库描述', blank=True, null=True)
    project_name = models.ForeignKey(to=ProjectInfo, on_delete=models.CASCADE,
                                     verbose_name='项目名称')
    
    status = models.BooleanField(choices=((1, '有效'), (0, '无效')), verbose_name='状态', default=1)
    create_time = models.DateTimeField(auto_now_add=True, null=False, blank=False, verbose_name='创建时间')
    update_time = models.DateTimeField(auto_now=True, verbose_name='更新时间')
    def __str__(self):
        return self.name

    class Meta:
        verbose_name = "历史对话库管理"
        verbose_name_plural = verbose_name


project_list = ProjectInfo.objects.all()
try:
    vector_db = MilvusDB([[c.name, c.real_name, map(int, c.dims.split(","))] for c in project_list if c.status==1])
except:
    vector_db = MilvusDB([])
        

class HistoryQA(models.Model):
    code = models.CharField(max_length=32, verbose_name='问答编码')
    question = models.TextField(verbose_name='问题')
    answer = models.TextField(verbose_name='回答')

    source = models.CharField(max_length=64, verbose_name='来源', blank=True, null=True)
    project_name = models.ForeignKey(to=ProjectInfo, on_delete=models.SET_DEFAULT, default=-1, verbose_name='项目名称名称')
    partition_name = models.ForeignKey(to=PartitionInfo, on_delete=models.SET_DEFAULT, default=-1, verbose_name='历史对话库名称')

    embedding_type = models.ForeignKey(to=EmbeddingType, on_delete=models.SET_DEFAULT, default=-1, verbose_name='特征向量类型')

    user_id = models.CharField(max_length=64, verbose_name='提问用户ID', blank=True, null=True)
    user = models.CharField(max_length=64, verbose_name='用户名', blank=True, null=True)
    session_id = models.CharField(max_length=32, verbose_name='对话轮次id', blank=True, null=True)
    meta_data = models.JSONField(max_length=256,  verbose_name='备注信息', blank=True, null=True)

    feedback = models.CharField(choices=(('1', '用户满意'), ('2', '用户不满意'), ('0', '未判定')), max_length=32, default='0', verbose_name='用户反馈')

    status = models.BooleanField(choices=((1, '有效'), (0, '无效')), verbose_name='状态', default=1)
    comment = models.CharField(max_length=256, verbose_name='备注', blank=True, null=True)
    create_time = models.DateTimeField(auto_now_add=True, null=False, blank=False, verbose_name='创建时间')
    update_time = models.DateTimeField(auto_now=True, verbose_name='更新时间')

    def __str__(self):
        return self.code

    class Meta:
        verbose_name = "历史对话记录"
        verbose_name_plural = verbose_name

    def save1(self, *args, **kwargs):
        super().save(*args, **kwargs)
        print("通过接口写入数据！")
 
    def save(self, *args, **kwargs):
        project_name = self.project_name.name
        partition_name = self.partition_name.name
        question = self.question
        embedding_type = self.embedding_type.name
        embedding_dim = embedding_name2dim_dict[embedding_type]
        # 修改数据
        if self.pk:
            # 在milvus中删除数据
            del_idx_dict = {embedding_dim: [self.code]}
            return_status = vector_db.delete_data(del_idx_dict, project_name)
            if return_status:
                print("修改数据")
                # 写入milvus
                model = choose_model(embedding_type)
                project_qs = ProjectInfo.objects.filter(name=project_name)
                if project_qs.exists():
                    project = project_qs[0]
                    dim_list = project.dims.split(",")
                    if project.status == 0:
                        project.status = 1
                        project.save()
                    if str(embedding_dim) not in dim_list:
                        dim_list.append(str(embedding_dim))
                        project.dims = ",".join(dim_list)
                        project.save()

                query_embedding = model(question)
                user = self.user if self.user else ""
                session_id = self.session_id if self.session_id else ""
                source = self.source if self.source else ""
                tmp_result_list = [
                    [self.code],
                    [question],
                    [query_embedding],
                    [user],
                    [session_id],
                    [embedding_type],
                    [source],
                    [datetime.now().strftime("%Y-%m-%d %H:%M:%S")]
                ]
                return_status, error_info = vector_db.insert_data(data=tmp_result_list, collection_name=project_name, partition_name=partition_name, embedding_dim=embedding_dim)
                if return_status:
                    super().save(*args, **kwargs)
        # 新增数据
        else:
            self.code = str(uuid.uuid4().hex[:8]) + str(int(round(time.time() * 1000)))
            print("增加数据")
            # 写入milvus
            model = choose_model(embedding_type)
            project_qs = ProjectInfo.objects.filter(name=project_name)
            if project_qs.exists():
                project = project_qs[0]
                dim_list = project.dims.split(",")
                if project.status == 0:
                    project.status = 1
                    project.save()
                if str(embedding_dim) not in dim_list:
                    dim_list.append(str(embedding_dim))
                    project.dims = ",".join(dim_list)
                    project.save()

            query_embedding = model(question)
            user = self.user if self.user else ""
            session_id = self.session_id if self.session_id else ""
            source = self.source if self.source else ""
            tmp_result_list = [
                [self.code],
                [question],
                [query_embedding],
                [user],
                [session_id],
                [embedding_type],
                [source],
                [datetime.now().strftime("%Y-%m-%d %H:%M:%S")]
            ]
            return_status, error_info = vector_db.insert_data(data=tmp_result_list, collection_name=project_name, partition_name=partition_name, embedding_dim=embedding_dim)
            if return_status:
                super().save(*args, **kwargs)
        print("后台直接写入数据！")
    

